[Bug] Reject invalid vocab-parallel CE labels - #4297

Summary

• Validate that every non-IGNORE_INDEX label is in [0, global_vocab_size) before vocab-parallel cross entropy starts its TP reductions.
• Use torch._assert_async so validation does not introduce an explicit device-to-host synchronization.
• Keep the existing forward and backward paths unchanged for valid labels and IGNORE_INDEX.

Why

For a label below zero other than IGNORE_INDEX, or a label greater than or equal to global_vocab_size, every TP rank treats the target as outside its local vocabulary shard. The forward path therefore clears every local target log-probability and all-reduces them to a zero loss. The backward path still retains the softmax term, however, and returns a nonzero logits gradient. That loss and gradient cannot describe the same mathematical function.

The previous version of this PR masked those gradients in backward. Although that restored forward/backward consistency, it silently treated malformed labels like ignored labels and diverged from regular cross entropy, which rejects out-of-range targets. Validating in forward is the smaller and safer contract: bad training data fails immediately instead of being hidden, and backward needs no special case.

Labels are replicated across the TP dimension, so every rank evaluates the same predicate. The check happens before the shard-local label transformation and before TP collectives.

Validation

• TP=2 Gloo repro: labels -1, vocab_size, and values above vocab_size are rejected on both ranks.
• TP=2 parity check: valid boundary labels and IGNORE_INDEX match torch.nn.functional.cross_entropy in forward and backward.
• py_compile and Pyrefly passed for torchtitan/components/loss.py.
• Pre-commit hooks passed for the changed file (Pyrefly was run separately).

Only the production file is included; no test files are committed.